package main

import (
	"container/heap"
	"fmt"
	"sync"
	"time"
)

// 任务结构体
type Task struct {
	value    interface{} // 数据
	priority int         // 优先级（数值越小，优先级越高）
}

// 优先级队列实现
type PriorityQueue []*Task

func (pq PriorityQueue) Len() int           { return len(pq) }
func (pq PriorityQueue) Less(i, j int) bool { return pq[i].priority < pq[j].priority }
func (pq PriorityQueue) Swap(i, j int)      { pq[i], pq[j] = pq[j], pq[i] }

// Push 插入新元素
func (pq *PriorityQueue) Push(x interface{}) {
	*pq = append(*pq, x.(*Task))
}

// Pop 移除优先级最高的元素
func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	*pq = old[0 : n-1]
	return item
}

// 数据总线结构体
type Bus struct {
	mu       sync.Mutex       // 保证线程安全
	queue    PriorityQueue    // 优先级队列
	notEmpty *sync.Cond       // 条件变量，用于通知消费者
	wg       sync.WaitGroup   // 等待组，用于等待消费者完成任务
	stopped  bool             // 标志生产是否停止
}

// 创建新的数据总线
func NewBus() *Bus {
	b := &Bus{
		queue: PriorityQueue{},
	}
	heap.Init(&b.queue) // 初始化优先级队列
	b.notEmpty = sync.NewCond(&b.mu)
	return b
}

// 异步生产者入栈
func (b *Bus) Produce(data interface{}, priority int) {
	b.mu.Lock()
	defer b.mu.Unlock()

	heap.Push(&b.queue, &Task{value: data, priority: priority})
	b.notEmpty.Signal() // 通知消费者有新任务
}

// 异步消费者出栈
func (b *Bus) Consume(consumerID int, processFunc func(data interface{})) {
	defer b.wg.Done()
	for {
		b.mu.Lock()
		for len(b.queue) == 0 && !b.stopped {
			b.notEmpty.Wait() // 等待生产者放入数据或总线停止
		}

		if len(b.queue) == 0 && b.stopped {
			b.mu.Unlock()
			return // 队列为空且总线已停止，退出
		}

		task := heap.Pop(&b.queue).(*Task) // 取出优先级最高的任务
		b.mu.Unlock()

		fmt.Printf("Consumer %d processing: %v (priority: %d)\n", consumerID, task.value, task.priority)
		processFunc(task.value)
	}
}

// 启动多个消费者
func (b *Bus) StartConsumers(consumerCount int, processFunc func(data interface{})) {
	for i := 0; i < consumerCount; i++ {
		b.wg.Add(1)
		go b.Consume(i+1, processFunc)
	}
}

// 停止总线（等待所有消费者完成任务）
func (b *Bus) Stop() {
	b.mu.Lock()
	b.stopped = true
	b.notEmpty.Broadcast() // 通知所有等待中的消费者退出
	b.mu.Unlock()
	b.wg.Wait()
}

func main() {
	// 创建一个优先级队列总线
	bus := NewBus()

	// 启动 3 个消费者
	bus.StartConsumers(3, func(data interface{}) {
		// 模拟耗时处理
		time.Sleep(500 * time.Millisecond)
		fmt.Printf("Processed: %v\n", data)
	})

	// 多个生产者入栈
	go func() {
		for i := 0; i < 10; i++ {
			priority := 10 - i // 优先级（数字越小越高）
			fmt.Printf("Produced: %d (priority: %d)\n", i, priority)
			bus.Produce(i, priority)
			time.Sleep(100 * time.Millisecond)
		}
	}()

	// 主线程等待生产结束
	time.Sleep(5 * time.Second)

	// 停止总线，等待所有消费者完成
	bus.Stop()
	fmt.Println("All tasks completed.")
}
